/*  Author: Kuljit S. Virk

    This file contains ALL the functions used in propagation of the
    Green function array. The basic flow is as follows.

    Let (t, tau) describe the lower triangular part of the 2-time
    plane. That is, t runs along points on the horizontal line, where 
    tau = 0, and tau > 0 is the vertical line reaching upto the
    diagonal.

    Let G(t,tau;k) be the G.F. matrix at times (t,tau) and BZ point k.

    The G.F. and all the related quantities (e.g. Self energy) are
    data structures where a single unit of this structure corresponds 
    to ONE value of t as follows:

    A linear array of ntau x nk many matrices. Each matrix is nband x nband 
    complex double matrix (the cx_mat type if armadillo is used). However typedef
    is used in order to use the alias complex_matrix so that it can be changed in
    one place to change it everywhere.

    The array for matrix_function is arranged as follows:

    +------------+------------+.....+---------------+-----------+....+---------------+
    | k=0, tau=1 | k=0, tau=2 |.....| k=0, tau=ntau | k=1 tau=0 |....| k=nk tau=ntau |
    +------------+------------+.....+---------------+-----------+....+---------------+


    From outside:
    1. propagator(...) is called
    2. propagator() iterates over the times along the diagonal
    and calls one_step() to integrate G(t,tau;k) from t to t + dt.
    3. one_step() does this for EACH k, and EACH tau > 0, by iterating
    over these two variables and calling integrator() at each
    point. 
    4. integrator() basically implements an implicit midpoint rule
    5. For tau = 0, a different integrator is called which propagates
    via the Schrodinger equation for the density matrix.
    5. Once a single step is propagated, the code returns to
    propagator where the solution is stored, and the next step is
    prepared.

*/

#include <propagator.h>
#include <selfenergy.h>

//=============================================================================
// F_current, F_next, F_self_energy are pointers to matrices of what dimensions
// Hamiltonian1 and 2 are pointers to matrices of what dimensions
void integrator(const parameters& params,
		const double half_dt,
		const complex_matrix& Hamiltonian1,
		const complex_matrix& Hamiltonian2,
		const complex_matrix& self_energy, 
		const complex_matrix& F_current,
		const complex_matrix& F_self_energy,
		complex_matrix& F_next)
{

  using namespace arma;
  int dim = params.num_bands;
  cx_mat Y1(dim,dim),Y2(dim,dim);
  cx_mat F(dim,dim), Fn(dim,dim), Fnext(dim,dim);
  cx_mat Id(dim,dim);

  MKL_Complex16 one_eye(0.0,1.0);
  MKL_Complex16 minus_one_eye(0.0,-1.0);

  //cout <<one_eye<<endl;

  Id.eye();

  int mesh_index1, mesh_index2,it,ik;

  Fn  = F_current;   
  Y1  = minus_one_eye*half_dt* (Hamiltonian1);
  Y1 += half_dt*(self_energy);

  Y2  = minus_one_eye*half_dt*(Hamiltonian2);
  Y2 += half_dt*(self_energy);

  //Fnext[ik] = F_current[ik] + (Y1+Y2)*F_current[ik] + F_self_energy;
  F   = Y1*Fn;
  F  += Fn;

  F  += F_self_energy;

  Y2  = Id - Y2;  
  Fn  = solve(Y2,F);

  //cout << Y1 << endl;
  //cout << Y2 << endl;
  //cout << real(Fn) <<endl;
  //printf("%14.4f\n",Fn(0));

  F_next = Fn;
}
//-----------------------------------------------------------------------------
void one_step(const parameters& params, 
	      const double half_dt,
	      const int itdiag,
	      const arma::vec&               Field,
	      const arma::ivec&              Deltakx,
	      const arma::mat&               EnergyBands,
	      const complex_matrix_function& Hamiltonian,
	      const complex_matrix_function& F_current,
	      complex_matrix_function&       F_self_energy, 
	      complex_matrix_function&       self_energy,
	      complex_matrix_function&       F_next)
{  
  using namespace std;
  using namespace arma;
  
  int iteration=0;
  bool converged = false;
  double err=1.0;
  int itau,it, index, mesh_index,ik,N;
  const complex_matrix_function F_temp;
  complex_matrix H1,H2;

  while (err > params.tolerance  && iteration < params.maxiter){
    //TODO: compute_self_energy(F_temp, F_self_energy,self_energy);
    //if (iteration == 0) 
    //  compute_self_energy(F_current, F_self_energy,self_energy)
    //else
    //  compute_self_energy(F_next, F_self_energy,self_energy)


    // Propagate all points tau > 0 by one time step
    // For each tau along the vertical line, starting one time step below the diagonal
    // compute the F_next.
    // tau(0) = 0. So at itau=0 we propagate from DM at time t(itdiag)
    // just to the right of the diagonal by one time step. This means
    // that the point at ntau falls of the grid and is replaced by the
    // propagation from what lies at ntau-1 at time t(itdiag). This is
    // easily implemented by simply sending in F_next pointer one step
    // ahead of all other pointers. 
    
    for (ik = 0; ik < params.num_mesh_points ; ik ++ )
      {
	for ( itau = 0; itau   < params.ntmem-1 ; itau++ )
	  {
	    it = max(itdiag - itau,0);  
	    N  = int(ik/params.num_mesh_points_x)*params.num_mesh_points_x;
	    mesh_index = min(ik + Deltakx(it), N + params.num_mesh_points_x-1);
	    index = itau + ik*params.ntmem;

	    H1 = Hamiltonian[ik];
	    H2 = Hamiltonian[mesh_index]; 
	    H1 *= Field(itdiag-1);
	    H2 *= Field(itdiag);
	    
	  for (int i=0; i < params.num_bands; i++)
	    {
	      H1(i,i) += EnergyBands(ik,i);
	      H2(i,i) += EnergyBands(ik,i);
	    }

	    integrator(params, half_dt,
		       H1,
		       H2,
		       self_energy[index],
		       F_current[index], 
		       F_self_energy[index], 
		       F_next[index + 1]);
	  }
      }
    iteration++;
  }
}
//=============================================================================
void integrator_diag(const parameters& params,
		     const double half_dt,
		     const complex_matrix& Hamiltonian1,
		     const complex_matrix& Hamiltonian2,
		     const complex_matrix& self_energy, 
		     const complex_matrix& F_current,
		     const complex_matrix& F_self_energy,
		     complex_matrix&       F_next)
{
  using namespace arma;

  int dim = params.num_bands;
  cx_mat Y1(dim,dim),Y2(dim,dim), U(dim,dim), Ubar(dim,dim);
  cx_mat F(dim,dim), Fn(dim,dim), Fnext(dim,dim);
  cx_mat Id(dim,dim);

  cx_vec eigval;
  cx_mat eigvec;

  MKL_Complex16 one_eye(0.0,1.0);
  MKL_Complex16 minus_one_eye(0.0,-1.0);

  //cout <<one_eye<<endl;

  Id.eye();

  int mesh_index1, mesh_index2,it,ik;

  Fn = F_current;

  Y1  = minus_one_eye*half_dt* (Hamiltonian1);
  Y1 += half_dt*(self_energy);

  Y2  = minus_one_eye*half_dt*(Hamiltonian2);
  Y2 += half_dt*(self_energy);


  // Approximate U = exp(1i H dt) = (I - 1i Hdt/2)/(I + 1i Hdt/2) 
  //Y1   = Id + Y1;
  //Y2   = Id - Y2;
  //U    = solve( Y2, Y1);

  // Since we always have 2-3 bands only, it is better to 
  // exponentiate directly
  eig_gen(eigval,eigvec,Hamiltonian1); 
  U = eigvec * diagmat(exp(minus_one_eye*eigval*half_dt)) * inv(eigvec);
  eig_gen(eigval,eigvec,Hamiltonian2);
  U = inv(eigvec) * U;
  U = eigvec * diagmat(exp(minus_one_eye*eigval*half_dt)) * U;


  Ubar = inv(U);
  F   = U*Fn*Ubar;


  // The scattering term 
  Fn  = F_self_energy;
  Fn  = U*Fn*Ubar;

  // Add the scattering term 
  F_next = F + Fn;
}

void one_step_diag(
		   const parameters&              params, 
		   const double                   half_dt,
		   const int                      itdiag,
		   const arma::vec&               Field,
		   const arma::ivec&              Deltakx,
		   const arma::mat&               EnergyBands,
		   const complex_matrix_function& Hamiltonian,
		   const complex_matrix_function& F_current, 
		   complex_matrix_function&       F_self_energy, 
		   complex_matrix_function&       self_energy,
		   complex_matrix_function&       F_next)
{
  using namespace arma;

  int iteration=0;
  bool converged = false;
  double err=1.0;
  int itau, index, mesh_index,ik,N;
  const cx_mat *F_temp;
  complex_matrix H1,H2;
  while (err > params.tolerance  && iteration < params.maxiter)
    {
      //TODO: compute_self_energy(F_temp, F_self_energy,self_energy);
      //if (iteration == 0) 
      //  compute_self_energy(F_current, F_self_energy,self_energy);
      //else
      //  compute_self_energy(F_temp, F_self_energy,self_energy);
      
      
      // Propagate all points tau > 0 by one time step
      // Here we fill the array at F_next[0], which is left untouched by
      // the integrator() function. Here we fill it by propagating from
      // F_current[0], which holds the DM at the present time.
      for (ik = 0; ik < params.num_mesh_points ; ik ++ )
	{
	  N          = int(ik/params.num_mesh_points_x)*params.num_mesh_points_x;
	  mesh_index = std::min(ik + Deltakx(itdiag), N + params.num_mesh_points_x-1);
	  index      = ik*params.ntmem;

	  H1 = Hamiltonian[ik];
	  H2 = Hamiltonian[mesh_index]; 
	  H1 *= Field[itdiag-1];
	  H2 *= Field[itdiag];
	  for (int i=0; i < params.num_bands; i++)
	    {
	      H1(i,i) += EnergyBands(ik,i);
	      H2(i,i) += EnergyBands(ik,i);
	    }

	  integrator_diag(params, half_dt,
			  H1,
			  H2,
			  self_energy[index],
			  F_current[index],
			  F_self_energy[index],
			  F_next[index]);

	  //cout <<"ik "<<ik<<"  E[ik] "<<EnergyBands[ik] <<endl;
	  //cout <<"F_current\n"<< F_current[index] <<endl;
	  //cout <<"F_next\n"<< F_next[index] <<endl;
	}
      iteration++;
    }
}
//=============================================================================

void propagation::propagator(const parameters&              params,
			     const arma::vec&               tdiag, 
			     const arma::vec&               tau, 
			     const arma::vec&               Efield,
			     const arma::mat&               EnergyBands,
			     const complex_matrix_function& IntHamiltonian,
			     const h5files_container&       h5files,
			     complex_matrix_function&       F_current)
{
  using namespace arma;

  std::cout<<"\npropagation::propagator():\n";
  double half_dt;

  // Construct the array of number of Dk imparted in the time interval
  ivec Deltakx(tdiag.n_elem);
  vec  Field(tdiag.n_elem);
  Deltakx.zeros();
  Field.zeros();
  double dA;
  for (int i=1; i < tdiag.n_elem; i++)
    {
      dA    = (Efield(i)+Efield(i-1))*(tdiag(i)-tdiag(i-1))/2;
      Deltakx(i)=int(dA/params.deltakx);
      Field(i) = Field(i-1) + dA;
    }
  Deltakx.save("Deltakx.dat",raw_ascii);
  Field.save("IntegratedField.dat",raw_ascii);
  Deltakx *= 0;
  // Allocation of workspace arrays
  cout << "Allocating workspace arrays" << endl;
  int nk     = params.num_mesh_points;
  int ntau   = params.ntmem;
  int nbands = params.num_bands;
  int ntot   = nk*ntau;

  // Work arrays: to be used in steppers
  complex_matrix_function F_next(ntot);
  complex_matrix_function F_self_energy(ntot);
  complex_matrix_function self_energy(ntot);

  int itdiag, itau;
  for ( int i = 0; i < ntot; i++){ 
    F_next[i].zeros(nbands,nbands);
    F_self_energy[i].zeros(nbands,nbands); 
    self_energy[i].zeros(nbands,nbands);
  }

  // Start the loop
  cout << "Loop start" << endl;

  complex_matrix_ptr F_temp;
  complex_matrix soldm(params.num_bands*params.num_bands, params.ntdiag*params.num_mesh_points);
  complex_matrix sol(params.ntmem,params.ntdiag);
  int iall;


  // Store the t=0 state
  itdiag=0;
  int itaustore;
  for (itaustore=0;itaustore<tau.n_elem;itaustore++)
    sol(itaustore,itdiag) = F_current[itaustore](0,0);

  for (int ik=0; ik<params.num_mesh_points; ik++)
    {
      soldm(0,itdiag+ik*params.ntdiag) = F_current[ik*params.ntmem](0,0);
      soldm(1,itdiag+ik*params.ntdiag) = F_current[ik*params.ntmem](1,1);
      soldm(2,itdiag+ik*params.ntdiag) = F_current[ik*params.ntmem](0,1);
      soldm(3,itdiag+ik*params.ntdiag) = F_current[ik*params.ntmem](1,0);
    }
  for ( itdiag = 1; itdiag < params.ntdiag; itdiag ++ ){
    if (itdiag < params.ntdiag-1)
      half_dt = (tdiag(itdiag+1)-tdiag(itdiag))/2;
    else
      half_dt = (tdiag(itdiag)-tdiag(itdiag-1))/2;

    printf("%10s %5d \n","itdiag",itdiag);
    one_step(params,half_dt,itdiag,Field,Deltakx,EnergyBands,IntHamiltonian,
	     F_current,F_self_energy,self_energy,F_next);
 
    one_step_diag(params,half_dt,itdiag,Field,Deltakx,EnergyBands,IntHamiltonian,
		  F_current,F_self_energy,self_energy,F_next);

    for (itaustore=0;itaustore<tau.n_elem;itaustore++)
      sol(itaustore,itdiag) = F_next[itaustore](0,0);

    for (int ik=0; ik<params.num_mesh_points; ik++)
      {
	soldm(0,itdiag+ik*params.ntdiag) = F_next[ik*params.ntmem](0,0);
	soldm(1,itdiag+ik*params.ntdiag) = F_next[ik*params.ntmem](1,1);
	soldm(2,itdiag+ik*params.ntdiag) = F_next[ik*params.ntmem](0,1);
	soldm(3,itdiag+ik*params.ntdiag) = F_next[ik*params.ntmem](1,0);
      }
    //cout << "itdiag, itmem = "<<itdiag<<" "<<itdiag%params.ntmem<<endl;
    io::dump(h5files[0],F_next,params,itdiag % params.ntmem);

    // Swap the current and next vectors for the next iteration. This way
    // we use the F_next as the F_current automatically, and overwrite
    // space allocated for F_current with the new solution. 
    F_current.swap(F_next);

    // Since F_next is recycling F_current, set F_next to zeros
    for (iall = 0; iall < ntot; iall++)
      F_next[iall].zeros();

  }

  mat r;
  r=real(sol).t();
  r.save("sol_real.dat",raw_ascii);
  r=imag(sol).t();
  r.save("sol_imag.dat",raw_ascii);
  tdiag.save("time.dat",raw_ascii);

  r=real(soldm).t();
  r.save("soldm_real.dat",raw_ascii);
  r=imag(soldm).t();
  r.save("soldm_imag.dat",raw_ascii);
  tdiag.save("time.dat",raw_ascii);


}

